from io import BytesIO

from flask import Blueprint, request
from flask import Response
from flask import jsonify
from rdkit import Chem
from rdkit.Chem import Draw

bp = Blueprint('rdkit', __name__, url_prefix='/rdkit')


@bp.post('/generate_image')
def generate_image():
    try:
        simle = request.form.get('simle')

        # 使用 RDKit 构建分子对象
        mol = Chem.MolFromSmiles(simle)
        if mol is None:
            return "Invalid SMILES representation", 400

        # 生成分子图片
        img = Draw.MolToImage(mol, size=(1000, 1000))  # 设置图片尺寸

        # 将图片保存到字节流中
        img_buffer = BytesIO()
        img.save(img_buffer, format='PNG')
        img_data = img_buffer.getvalue()

        # 返回图片作为响应
        return Response(img_data, content_type='image/png')

    except Exception as e:
        res = {
            "code": 500,
            "msg": str(e),
        }
        print(e)
        return jsonify(res)
